package com.agentai.model.controller;


import com.agentai.framework.common.utils.Result;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.Media;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.net.URL;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

@Slf4j
@RequestMapping("/deep")
@RestController
public class ChatDeepSeekController {

    @Resource
    private OllamaChatModel ollamaChatModel;
    @Resource
    private OpenAiChatModel openAiChatModel;
    @GetMapping("/deep/seek/message")
    public Result<ChatResponse> message(@RequestParam("message") String message) {
        message = "请使用中文简体回答：" + message;
        Prompt prompt = new Prompt(new UserMessage(message));
        ChatResponse msg = openAiChatModel.call(prompt);
        return Result.ok(msg);
    }
    /**
     * 流式调用
     */
    @GetMapping("/generateStream")
    public SseEmitter streamChat(@RequestParam String message) {
        // 创建 SSE 发射器，设置超时时间（例如 1 分钟）
        SseEmitter emitter = new SseEmitter(60_000L);
        // 创建 Prompt 对象
        Prompt prompt = new Prompt(new UserMessage(message));
        // 订阅流式响应
        // 完成处理
        // 异常处理
        openAiChatModel.stream(prompt).subscribe(response -> {
                    try {
                        String content = response.getResult().getOutput().getText();
                        System.out.print(content);
                        // 发送 SSE 事件
                        emitter.send(SseEmitter.event()
                                .data(content)
                                .id(String.valueOf(System.currentTimeMillis()))
                                .build());
                    } catch (Exception e) {
                        emitter.completeWithError(e);
                    }
                },
                emitter::completeWithError,
                emitter::complete
        );
        // 处理客户端断开连接
        emitter.onCompletion(() -> {
            // 可在此处释放资源
            System.out.println("SSE connection completed");
        });
        emitter.onTimeout(() -> {
            emitter.complete();
            System.out.println("SSE connection timed out");
        });
        return emitter;
    }


    private Flux<ChatResponse> print(OllamaOptions options, Prompt prompt) {

        OllamaApi api = new OllamaApi();
        OllamaChatModel test = OllamaChatModel.builder().ollamaApi(api).defaultOptions(options).build();

        Flux<ChatResponse> chatResponseFlux = test.stream(prompt);
        chatResponseFlux
                .doOnNext(response -> {
                    for (Generation result : response.getResults()) {
                        log.info("Received response: {}", result.getOutput().getText());
                    }
                })
                .subscribe();

        return chatResponseFlux;
    }

    public static String convert(String imagePath) throws Exception {
        // 1. 读取图片文件
        BufferedImage image = ImageIO.read(new File(imagePath));

        // 2. 创建字节输出流
        ByteArrayOutputStream baos = new ByteArrayOutputStream();

        // 3. 获取文件扩展名（自动识别格式）
        String formatName = imagePath.substring(imagePath.lastIndexOf(".") + 1).toLowerCase();

        // 4. 将图像写入字节流
        if (!ImageIO.write(image, formatName, baos)) {
            throw new IllegalArgumentException("不支持的图片格式: " + formatName);
        }

        // 5. 转换为Base64字符串
        return Base64.getEncoder().encodeToString(baos.toByteArray());
    }

}
